from Crypto.Cipher import AES
from Crypto.Hash import MD5, SHA256
import StringIO
import pickle

import common

def PAD(s, multiple, pad_char = "\x00"):
    pad_length = multiple - (len(s) % multiple)
    return s + (pad_char * pad_length)

def AES_PAD(data):
    return PAD(data, AES.block_size)

def NEW_AES(*args):
    key = SHA256.new()
    iv = MD5.new()
    for arg in args:
        key.update(arg)
        iv.update(arg)
    iv.update(".iv")
    return AES.new(key.digest(), IV=iv.digest())

def AES_ENCRYPTOR(*args):
    return NEW_AES(*args)

def NAME_ENCRYPTOR(password):
    return AES_ENCRYPTOR(password, ".presetName")

def OBJ_ENCRYPTOR(name, password):
    return AES_ENCRYPTOR(password, name, ".presetObj")


class PresetFile:
    def __init__(self, filename):
        self.filename = filename
    def save(self, preset_dict, password):
        output_dict = {}
        for preset_name, preset_obj in preset_dict.iteritems():
            encrypted_name = NAME_ENCRYPTOR(password).encrypt(AES_PAD(preset_name))
            encrypted_obj = OBJ_ENCRYPTOR(preset_name, password).encrypt(AES_PAD(pickle.dumps(preset_obj)))
            assert encrypted_name not in output_dict, "Internal error occurred"
            output_dict[encrypted_name] = encrypted_obj
        output_file = open(self.filename, "wb")
        pickle.dump(output_dict, output_file)
        output_file.close() 
        
    def load(self, password):
        try:
            result = {}
            input_file = open(self.filename, "rb")
            encrypted_dict = pickle.load(input_file)
            input_file.close()
            for encrypted_name, encrypted_obj in encrypted_dict.iteritems():
                preset_name = NAME_ENCRYPTOR(password).decrypt(encrypted_name).replace("\x00", "")
                preset_obj = OBJ_ENCRYPTOR(preset_name, password).decrypt(encrypted_obj)
                preset_obj = pickle.loads(preset_obj)
                result[preset_name] = preset_obj
            return result
        except Exception, e:
            if common.VERBOSE:
        	      print e
            raise common.PresetFileException("Unable to load preset file. Invalid password?")

if __name__ == "__main__":
    pfile = PresetFile(r"c:\temp\file.mkp")
    d = {"bla":"rotem"}
    pfile.save(d, "rrr")
    d2 = pfile.load('rrr')
    if d2 != d:
        raise Exception("ERR!!!")
